import re
import pandas as pd
from fuzzywuzzy import fuzz
from Agent.Agent import Agent
from llm import llm, super_eval
from kernel import CodeKernel
import json

data = pd.read_excel('./字段释义.xlsx')
data = data.to_dict('records')
text_list = []


dic = {'device_13_2_meter_1302': '二号舵桨转舵A电压电流及功率测量表',
       'Port2_ksbg_1': '柴油发电机组运行参数表', 'Port2_ksbg_2': '三号柴油发电机组有功功率测量表',
       'Port1_ksbg_5': '左舵桨主开关电流测量表', 'device_1_2_meter_102': '一号舵桨转舵A电压电流及功率测量表',
       'Port1_ksbg_2': '停泊_应急发电机组运行参数表',
       'device_13_11_meter_1311': '折臂吊车液压系统电压电流及功率测量表',
       'Port4_ksbg_7': '二号推进变频器功率测量表', 'Port2_ksbg_4': '推进系统主开关电流测量表',
       'Ajia_plc_1': 'A架角度和启动柜电压电流测量表',
       'device_13_14_meter_1314': '二号门架主液压泵电压电流及功率测量表',
       'Port4_ksbg_9': '二号舵桨方位及转速测量表', 'Port3_ksbg_8': '一号推进变频器功率测量表',
       'Port1_ksbg_4': '推进系统及柴油发电机组状态测量表', 'Port3_ksbg_10': '一号舵桨方位及转速测量表',
       'Port1_ksbg_3': '一号和二号柴油发电机组有功功率及运行状态测量表',
       'device_1_3_meter_103': '一号舵桨转舵B电压电流及功率测量表',
       'device_1_15_meter_115': '绞车变频器电压电流及功率测量表',
       'Port2_ksbg_3': '三号和四号柴油发电机组有功功率及运行状态测量表',
       'device_1_5_meter_105': '一号门架主液压泵电压电流及功率测量表',
       'Jiaoche_plc_1': '绞车A、B、C放缆长度、速度及张力测量表',
       'Port1_ksbg_1': '一号和二号柴油发电机组运行参数及燃油消耗率测量表',
       'Port3_ksbg_9': '艏侧推螺旋桨螺距及功率测量表',
       'device_13_3_meter_1303': '二号舵桨转舵B电压电流及功率测量表',
       'Port4_ksbg_8': '可伸缩推螺距、方位及功率测量表'}

for i in data:
    try:
        i['所属表名'] = dic[i['所属表名']]
    except:
        ...

for i in data:
    if str(i['单位']) != 'nan':
        text_list.append(str(i['字段含义']) + str(i['单位']) + str(i['归属']))
    else:
        text_list.append(str(i['字段含义']) + str(i['归属']))

with open('./sample/function.json', encoding='utf-8') as f:

    code_text_list = json.load(f)


def search_code(question):
    score_list = [(s, fuzz.ratio(question, s)) for s in code_text_list]

    top_10 = sorted(score_list, key=lambda x: x[1], reverse=True)[:6]

    return '\n\n'.join([i[0] for i in top_10])


prompt = r'''
以下是参考实现：
```python
[code]
```
以下是数据表格：
---
[data]
---
备注：所有表格都有`时间`列
你会基于用户的问题，编写python代码解决问题，
你优先使用pandas内部计算。
代码写在
```python
```
内。
其中，一些公式为：

柴油发电机组的运⾏可根据相应的【额定转速运⾏字段】（0为关，1为开）。其他的设备运⾏全部根
据功率判断，功率>0则判定为运⾏，运⾏时⻓可以直接计算满⾜条件的⾏数，如某段时间内有⼗条满
⾜运⾏条件的数据，则运⾏时⻓为10分钟。
3、关于查询XX时间的设备参数问题
如：请确定2024年8⽉19⽇下午A架的第⼀次开机时间（格式：HH:MM），并报告开机时1号柴油发
电机组的功率（单位：kW，输出为整数）。
先查询出动作发⽣的具体时间（HH:MM:SS），然后查询1号柴油发电机组的数据中时间离这个动作发
⽣时间最近的功率数据。⽐如查询出来时间是13:01:01，此时对应1号柴油发电机组的数据有13:00:50
和13:01:50，则取13:00:50的⾏所对应的值。
如果题⽬中有出现如【发⽣后10分钟内】类似字眼，还是以这道题为例⼦，查询出来13:01:01，则发
⽣后⼗分钟内是13:01:01~13:11:01的闭区间。
4、关于功率/做功/燃油消耗
查询功率值从【有功功率】【功率反馈】⾥找。
如果询问设备的总功率，需要将各个部分的功率求和。⽐如发电机组总功率则是将1~4加加起来，主推
进器则是将⼀号、⼆号推进器功率加起来。
做功和燃油消耗同理。做功/燃油消耗量的具体计算公式沿⽤初赛给出的补充资料，该⽂档下也会附
上。
5、其他
 - 回答关于是否类问题如果没有明确要求回答格式，使⽤Y或N；
 - 关于参数查询，如果没有写明，可以按照查询的值直接给出，例如可伸缩推方位反馈范围为-2000%-2000%都是合理的，除非有备注，否则无需处理。
附录
时⻓/能耗相关问题解释
1、关于【运⾏时⻓/开机时间/实际运⾏时⻓】
a) ⽆特殊说明的情况下，运⾏时⻓=开机时⻓。
b) 实际运⾏时⻓是电流或其他相关参数⼤于待机状态的时⻓。以A架为例，即两边电流均⼤于0时判定
为实际运⾏。
2、关于计算，因为数据有缺失，尽量考虑平均值*时间的方法计算，
以能耗计算为例，使⽤下⾯的计算⽅式：

假设：
- 功率列的数据范围为 \(P_1, P_2, P_3, \ldots, P_n\)。
- 时间列的数据范围为 \(T_1, T_2, T_3, \ldots, T_n\)。

则总做功 \(W\) 的公式为：

\[ W = \sum_{i=1}^{n-1} P_i \times \frac{(T_{i+1} - T_i)}{3600} \]

公式说明：
1. \(P_i\)：第 \(i\) 行的功率值。
2. \(T_{i+1} - T_i\)：第 \(i\) 行与第 \(i+1\) 行之间的时间间隔（单位为秒）。
3. \(\frac{(T_{i+1} - T_i)}{3600}\)：将时间间隔从秒转换为小时。
4. \(P_i \times \frac{(T_{i+1} - T_i)}{3600}\)：第 \(i\) 行的功率乘以时间间隔，得到该行的做功。
5. \(\sum_{i=1}^{n-1}\)：对所有行的做功结果累加，得到总做功。

a）某时间跨度内的能耗计算：1）筛选出该时间段的数据（闭区间）；2）使⽤公式进⾏计算
b）两个动作期间的能耗计算：1）判断动作发⽣时间（XX:XX:XX）(HH:MM:SS)；2）将两个动作的时
间作为开始和结束点，进⾏a）的计算。
燃油消耗的计算同理，只不过把功率字段换成燃油消耗率的字段。
3、关于能耗归属
推进系统：主推进器（⼀号/⼆号推进器）/艏侧推（也叫艏推或者侧推）/（可）伸缩推
甲板机械系统：A架（⼀号/⼆号⻔架）/绞⻋（A/B/C）/折臂吊⻋
柴油发电机：1~4号，除⾮指明需要计算应急柴油发电机，否则不考虑

```特别注意，计算能耗时，你需要找到所有的运作机器，并只取开机时间作为运行时常。```

我的数据都存放在 `./data2/表格名称.csv` 下，你可以读取，所有表格都有`时间`列，格式为`2024-05-20 07:02:00`，每分钟一条数据(可能会缺失)
请尽可能分步编码，使用简单但是有逻辑的方法书写代码，你在第一次编码时应该print(df.head(5)),以确定数据使用的是否正确。
确定数据后进行编码。
一般步骤为：
1.读取文件并按照时间过滤，只选择有用的列，删除无用列
2.按照时间合并文件，并按照时间排序
3.打印合并后的文件，
4.编写计算规则，
5.print 计算结果。
示例：
假设柴油的密度为0.85kg/L，柴油热值为42.6MJ/kg。请根据提供的1~4号柴油发电机的燃油消耗量，计算2024年7月16日00:00:00至2024年7月20日00:00:00期间的理论发电量（单位转换为kWh，结果保留两位小数，四舍五入）'
```python
import pandas as pd

# 定义时间范围
start_time = pd.to_datetime('2024-07-16 00:00:00')
end_time = pd.to_datetime('2024-07-20 00:00:00')

# 读取并处理Port1_ksbg_1数据
df_port1_ksbg1 = pd.read_csv('./data2/Port1_ksbg_1.csv', parse_dates=['时间'])
df1 = df_port1_ksbg1[['时间', '一号柴油发电机组燃油消耗率', '二号柴油发电机组燃油消耗率']].copy()
df1 = df1[(df1['时间'] >= start_time) & (df1['时间'] <= end_time)]

# 读取并处理Port2_ksbg_1数据
df_port2_ksbg1 = pd.read_csv('./data2/Port2_ksbg_1.csv', parse_dates=['时间'])
df2 = df_port2_ksbg1[['时间', '三号柴油发电机组燃油消耗率', '四号柴油发电机组燃油消耗率']].copy()
df2 = df2[(df2['时间'] >= start_time) & (df2['时间'] <= end_time)]

def calculate_fuel(df, column):
    """计算指定列的总燃油消耗量"""
    if df.empty:
        return 0.0
    df = df.sort_values('时间').reset_index(drop=True)
    # 计算相邻时间间隔（小时）
    df['delta'] = df['时间'].diff().dt.total_seconds().shift(-1).fillna(0) / 3600
    # 计算各时段燃油消耗量
    df['fuel'] = df[column] * df['delta']
    return df['fuel'].sum()

# 计算各机组燃油消耗量
fuel_1 = calculate_fuel(df1, '一号柴油发电机组燃油消耗率')
fuel_2 = calculate_fuel(df1, '二号柴油发电机组燃油消耗率')
fuel_3 = calculate_fuel(df2, '三号柴油发电机组燃油消耗率')
fuel_4 = calculate_fuel(df2, '四号柴油发电机组燃油消耗率')

# 总燃油消耗量（升）
total_fuel = fuel_1 + fuel_2 + fuel_3 + fuel_4

# 计算理论发电量（kWh）
density = 0.85  # kg/L
heating_value = 42.6  # MJ/kg
total_energy = (total_fuel * density * heating_value) / 3.6  # 1kWh=3.6MJ

# 结果格式化
print(f"{total_energy:.2f}kWh")
```
'''


class CodeAgent(Agent):

    def __init__(self, name,
                 description,
                 system_prompt,
                 agent_type,
                 llm_engine,
                 agents=None,
                 rag=None, sample=None):
        super().__init__(name,
                         description,
                         system_prompt,
                         agent_type,
                         llm_engine,
                         agents,
                         rag)

        self.sample = sample
        self.kernel = CodeKernel()
        self.kernel.start()
        self.kernel.run('''
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "last_expr"
# 简化异常显示
%xmode Minimal
import warnings
warnings.filterwarnings("ignore")
import pandas as pd
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 20)
pd.set_option('display.float_format', lambda x: str(x))
from typing import Optional, List, Dict, Any
import os
from datetime import datetime
''')
    def get_column_data(self, question):
        # 计算相似度
        scores = [(index, fuzz.ratio(text, question)) for index, text in enumerate(text_list)]
        # 排序并获取前50个最相似的元素的索引
        sorted_scores = sorted(scores, key=lambda x: x[1], reverse=True)
        top_100_indices = [index for index, score in sorted_scores[:77] if score > 2]
        # 从 data 中提取相应的记录
        top_50_records = [data[index] for index in range(len(data)) if index in top_100_indices]
        res = pd.DataFrame(top_50_records)
        res['字段名'] = res['字段含义']
        del res['字段含义']
        return res.to_markdown(index=False)

    def build_prompt(self, question):
        code_recall = search_code(question)
        self.messages = [
            {"role": "system",
             "content": self.system_prompt.replace('[data]', self.get_column_data(question)).replace('[code]',code_recall)},
        ]
        recall_sample = self.sample.get_sample(question, 2)
        for qa in recall_sample:
            self.messages.extend(qa)

        self.messages.append({"role": "user", "content": question + "\n请思考后编码回答问题。"})

    def run(self, question):


        self.build_prompt(question)

        for try_num in range(7):
            if try_num <= 3:
                msg = '如果答案为0 你需要多确认1次，无需反复确认'
            else:
                msg = '如果答案为0,也是合理的，请尽快修改代码或者回答问题。'
            answer = llm(self.messages)
            if '```python' in answer:
                code_res = ''
                code_list = re.findall(r'```python(.*?)```', answer, re.DOTALL)
                for code in code_list:
                    code_res += self.kernel.run(code)

                self.messages.append({"role": "user", "content": f"""
代码执行结果为：
---
{code_res}
---
如果答案不合理，请修改代码，忽略警告。
如果是没有找到数据列，从数据表格确认列名和所在表格，禁止假设。
{msg}
如果运行结果可以回答问题，且答案合理，则回答问题:{question}
"""})
            else:
                with open('./sample/code5.jsonl', 'a', encoding='utf-8') as fw:
                    fw.write(
                        json.dumps({"question": question, "messages": self.messages[1:]}, ensure_ascii=False)
                    )
                    fw.write('\n')

                return answer

        return '题目太难，请简化'

from get_sample import get_code_sample

CodeAgentConfig = {"name": 'xxx',
                   "description": 'xxx',
                   'system_prompt': prompt,
                   'agent_type': 'plan-action',
                   'llm_engine': llm,
                   'sample': get_code_sample(),
                   'agents': None}

if __name__ == '__main__':
    print(search_code(' 计算所有设备（折臂吊车、一号门架、二号门架、绞车）在指定时间范围内的总能耗'))
